import os
import json
import requests
from typing import List, Dict

# 星火AI接口地址
url = 'https://spark-api-open.xf-yun.com/v1/chat/completions'
# APIPassword, 本地开发可从环境变量改为静态字符串
api_token = os.getenv('SPARKAI_PASSWORD')
# 模拟用户的问题
question = '请介绍以下datawhale这个开源社区'

def chat(messages:List[Dict]):
    # 构建请求头
    headers = {
        'Content-Type': 'application/json',
        'Authorization': 'Bearer ' + api_token,
    }
    # 构建请求体
    req_body = {
        # 固定参数
        "model": "4.0Ultra",
        # 用户消息
        "messages": messages,
        # ----- 下面是可选参数 ----- #
        # 流式输出
        "stream": True
    }
    # 发起请求，注意stream与请求体要一只
    response = requests.post(url=url,json= req_body,headers= headers,stream= True)
    full_response = '' # 初始化模拟结果
    isFirstContent = True  # 首帧标识
    # 处理响应, 遍历http响应流
    for chunks in response.iter_lines():
        if (chunks and '[DONE]' not in str(chunks)):
            data_org = chunks[6:]
            chunk = json.loads(data_org)
            text = chunk['choices'][0]['delta']
            # 判断最终结果状态并输出
            content = text["content"]
            print(content)
            full_response += content
    return full_response


messages = [
    {"role": "user", "content": question}
]
chat(messages)